import math

try :    
    cmds.delete( dellist )
    dellist = []
except :
    dellist = []

def distance(a,b):
    posA = cmds.xform( a , q=1 , rp=1 )
    posB = cmds.xform( b , q=1 , rp=1 )
    tot = 0
    for i in range(3):
        tot +=(posA[i]-posB[i])**2
    return math.sqrt(tot)
    
mode = 1
radL = 0
radR = 0

def antwalk( ant ):
    global radL
    global radR
    ct = cmds.currentTime( q=1 )
    
    main = [n for n in cmds.listRelatives(ant) if 'main' in n]
    leg_grp = [n for n in cmds.listRelatives( ant ) if 'leg' in n]
    leg_ctrls = [n for n in cmds.listRelatives(leg_grp[0]) ]
    
    # initial
    if ct == 1 :
        map( lambda s :cmds.setAttr( 'ant_rtft_ctrl.t' + s , 0 ) , 'xy' )
        map( lambda s :cmds.setAttr( 'ant_rtmd_ctrl.t' + s , 0 ) , 'xy' )
        map( lambda s :cmds.setAttr( 'ant_rtbk_ctrl.t' + s , 0 ) , 'xy' )
        
        map( lambda s :cmds.setAttr( 'ant_lfft_ctrl.t' + s , 0 ) , 'xy' )
        map( lambda s :cmds.setAttr( 'ant_lfmd_ctrl.t' + s , 0 ) , 'xy' )
        map( lambda s :cmds.setAttr( 'ant_lfbk_ctrl.t' + s , 0 ) , 'xy' )        
        
    mainX = cmds.getAttr( 'ant_main.tx' )
    mainX1 = cmds.getAttr( 'ant_main.tx' , t = ct-1 )    
    dist2Main = mainX - mainX1

    upSpeed = 5.0
    liftH = .15
    
    ## Front

    if not dist2Main and not max( 0 , math.sin( radR ) * liftH ) :
        pass
    else:
        if cmds.getAttr('ant_lfft_ctrl.ty') == 0 :
            cmds.setAttr( 'ant_rtft_ctrl.ty' , max( 0 , math.sin( radR ) * liftH ) )       
            radR = ( radR + math.pi / upSpeed )        

    if cmds.getAttr('ant_rtft_ctrl.ty') != 0 :
        cmds.setAttr( 'ant_rtft_ctrl.tx' , mainX )        
    
    if not dist2Main and not max( 0 , math.sin( radL ) * liftH ) :
        pass
    else:
        if cmds.getAttr('ant_rtft_ctrl.ty') == 0 :
            cmds.setAttr( 'ant_lfft_ctrl.ty' , max( 0 , math.sin( radL ) * liftH ) )        
            radL = ( radL + math.pi / upSpeed )
            
    if cmds.getAttr('ant_lfft_ctrl.ty') != 0 :
        cmds.setAttr( 'ant_lfft_ctrl.tx' , mainX )
        
        
    ## Middle
    if not dist2Main and not max( 0 , math.sin( radR ) * liftH ) :
        pass
    else:
        if cmds.getAttr('ant_lfmd_ctrl.ty') == 0 :
            cmds.setAttr( 'ant_rtmd_ctrl.ty' , max( 0 , math.sin( radR ) * liftH ) )        
            radR = ( radR + math.pi / upSpeed )
                
    if cmds.getAttr('ant_rtmd_ctrl.ty') != 0 :
        cmds.setAttr( 'ant_rtmd_ctrl.tx' , mainX )        
    
    if not dist2Main and not max( 0 , math.sin( radL ) * liftH ) :
        pass
    else:
        if cmds.getAttr('ant_rtmd_ctrl.ty') == 0 :        
            cmds.setAttr( 'ant_lfmd_ctrl.ty' , max( 0 , math.sin( radL ) * liftH ) )        
            radL = ( radL + math.pi / upSpeed )
                
        
    if cmds.getAttr('ant_lfmd_ctrl.ty') != 0 :
        cmds.setAttr( 'ant_lfmd_ctrl.tx' , mainX )
        
        
        
    ## Back
    if not dist2Main and not max( 0 , math.sin( radR ) * liftH ) :
        pass
    else: 
        if cmds.getAttr('ant_lfbk_ctrl.ty') == 0 :
            cmds.setAttr( 'ant_rtbk_ctrl.ty' , max( 0 , math.sin( radR ) *  liftH ) )
            radR = ( radR + math.pi / upSpeed ) 
    
        
    if cmds.getAttr('ant_rtbk_ctrl.ty') != 0 :
        cmds.setAttr( 'ant_rtbk_ctrl.tx' , mainX )        
    
    
    if  dist2Main == 0  and  max( 0 , math.sin( radL ) * liftH ) ==0 :
        pass
    else:
        if cmds.getAttr('ant_rtbk_ctrl.ty') == 0 :        
            cmds.setAttr( 'ant_lfbk_ctrl.ty' , max( 0 , math.sin( radL ) * liftH ) )
            radL = ( radL + math.pi / upSpeed )        
        
    if cmds.getAttr('ant_lfbk_ctrl.ty') != 0 :
        cmds.setAttr( 'ant_lfbk_ctrl.tx' , mainX )



    if  dist2Main == 0  :
        map( lambda s :cmds.setAttr( 'ant_'+s+'ft_ctrl.ty' , 0 ) , [ 'rt' ,'lf' ] )
        map( lambda s :cmds.setAttr( 'ant_'+s+'md_ctrl.ty' , 0 ) , [ 'rt' ,'lf' ] )
        map( lambda s :cmds.setAttr( 'ant_'+s+'bk_ctrl.ty' , 0 ) , [ 'rt' ,'lf' ] )



if mode == 1:
    for x in range( cmds.playbackOptions(q=1,aet=1)):
        cmds.currentTime(x+1)
        antwalk()
        map( lambda s :cmds.setKeyframe( 'ant_'+s+'ft_ctrl' ,at='translateX' ) , ['lf','rt'] )
        map( lambda s :cmds.setKeyframe( 'ant_'+s+'md_ctrl' ,at='translateX' ) , ['lf','rt'] )
        map( lambda s :cmds.setKeyframe( 'ant_'+s+'bk_ctrl' ,at='translateX' ) , ['lf','rt'] )
        
        map( lambda s :cmds.setKeyframe( 'ant_'+s+'ft_ctrl' ,at='translateY' ) , ['lf','rt'] )
        map( lambda s :cmds.setKeyframe( 'ant_'+s+'md_ctrl' ,at='translateY' ) , ['lf','rt'] )
        map( lambda s :cmds.setKeyframe( 'ant_'+s+'bk_ctrl' ,at='translateY' ) , ['lf','rt'] )

else :
    exp = cmds.expression( s='python("antwalk()")' , ae = 1 )
    dellist.append( exp )




